import sys
import os
from pathlib import Path
from datetime import datetime
import json
import glob


sys.path.insert(0, os.path.dirname(__file__))


from mp3_json_enhanced import transcribe_audio_to_json, ENHANCED_MODE_AVAILABLE, HUBERT_MODE_AVAILABLE
from translation_module import translate_text
from tts_module import text_to_speech
from qwen_polish_module import polish_text


XFYUN_TOKEN_LIMIT = 5000
QWEN_TOKEN_LIMIT = 5000
HISTORY_DIR = Path(__file__).parent.parent / "数据" / "翻译历史记录"
QWEN_HISTORY_DIR = Path(__file__).parent / "ai调用" / "使用历史记录"


AVAILABLE_WHISPER_MODELS = {
    'tiny': {'size': '39MB', 'speed': '最快', 'accuracy': '一般', 'description': '快速测试'},
    'base': {'size': '74MB', 'speed': '快', 'accuracy': '较好', 'description': '日常使用'},
    'small': {'size': '244MB', 'speed': '中等', 'accuracy': '好', 'description': '平衡推荐'},
    'medium': {'size': '769MB', 'speed': '慢', 'accuracy': '很好', 'description': '高质量'},
    'large': {'size': '1550MB', 'speed': '最慢', 'accuracy': '最好', 'description': '最高精度'}
}


def get_total_chars_used():
    if not HISTORY_DIR.exists():
        return 0

    total_chars = 0
    json_files = glob.glob(str(HISTORY_DIR / "record_*.json"))

    for json_file in json_files:
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                total_chars += data.get("字符使用情况", {}).get("总计", 0)
        except:
            continue

    return total_chars

def get_total_qwen_tokens_used():
    if not QWEN_HISTORY_DIR.exists():
        return 0

    total_tokens = 0
    json_files = glob.glob(str(QWEN_HISTORY_DIR / "chat_*.json"))

    for json_file in json_files:
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                total_tokens += data.get("Token使用情况", {}).get("总计", 0)
        except:
            continue

    return total_tokens

def save_translation_history(data):
    HISTORY_DIR.mkdir(exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"record_{timestamp}.json"
    filepath = HISTORY_DIR / filename

    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    return filepath


def main():
    print("=" * 70)
    print("🎙️  音频翻译系统")
    print("=" * 70)


    chars_used = get_total_chars_used()
    chars_remaining = XFYUN_TOKEN_LIMIT - chars_used
    chars_percentage = (chars_used / XFYUN_TOKEN_LIMIT) * 100

    print(f"\n📊 讯飞API字符使用情况：{chars_used}/{XFYUN_TOKEN_LIMIT} ({chars_percentage:.1f}%)")
    print(f"   剩余可用：{chars_remaining} 字符")


    qwen_tokens_used = get_total_qwen_tokens_used()
    qwen_remaining = QWEN_TOKEN_LIMIT - qwen_tokens_used
    qwen_percentage = (qwen_tokens_used / QWEN_TOKEN_LIMIT) * 100

    print(f"\n📊 通义千问Token使用情况：{qwen_tokens_used}/{QWEN_TOKEN_LIMIT} ({qwen_percentage:.1f}%)")
    print(f"   剩余可用：{qwen_remaining} tokens")


    if chars_used >= XFYUN_TOKEN_LIMIT:
        print(f"\n❌ 讯飞API字符使用量已达到限制 ({XFYUN_TOKEN_LIMIT})！")
        print("为了控制成本，服务已暂停。")
        return

    if qwen_tokens_used >= QWEN_TOKEN_LIMIT:
        print(f"\n❌ 通义千问Token使用量已达到限制 ({QWEN_TOKEN_LIMIT})！")
        print("为了控制成本，服务已暂停。")
        return


    if chars_remaining < 500 or qwen_remaining < 500:
        print(f"\n⚠️  警告：剩余额度不足！")
        confirm = input("是否继续？(y/n)：").strip().lower()
        if confirm != 'y':
            print("已取消操作。")
            return

    print("\n" + "=" * 70)


    print("\n【步骤 1/5】选择翻译方向")
    print("-" * 70)
    print("1. 中文 → 法语（源语言：中文）")
    print("2. 法语 → 中文（源语言：法语）")

    while True:
        choice = input("请选择（1或2）：").strip()
        if choice == '1':
            direction = 'chinese_to_french'
            target_lang = 'french'
            source_lang = 'zh'
            print("✓ 已选择：中文 → 法语\n")
            break
        elif choice == '2':
            direction = 'french_to_chinese'
            target_lang = 'chinese'
            source_lang = 'fr'
            print("✓ 已选择：法语 → 中文\n")
            break
        else:
            print("❌ 无效的选择\n")


    print("=" * 70)
    print(f"\n【步骤 2/5】选择音频文件（源语言：{'中文' if source_lang=='zh' else '法语'}）")
    print("-" * 70)

    while True:
        audio_path = input("请输入音频文件路径（或输入文件名在音频仓库中查找）：").strip()
        if not audio_path:
            print("❌ 文件路径不能为空\n")
            continue

        audio_file = Path(audio_path)
        if not audio_file.is_absolute():
            audio_dir = Path(__file__).parent.parent / "音频仓库"
            audio_file = audio_dir / audio_path

        if audio_file.exists():
            print(f"✓ 找到音频文件：{audio_file}\n")
            break
        else:
            print(f"❌ 文件不存在：{audio_file}\n")


    print("\n可用的 Whisper 模型：")
    print("-" * 70)
    for model_name, info in AVAILABLE_WHISPER_MODELS.items():
        print(f"  {model_name:8} - {info['size']:8} | 速度: {info['speed']:4} | 精度: {info['accuracy']:4} | {info['description']}")
    print("-" * 70)

    VALID_MODELS = list(AVAILABLE_WHISPER_MODELS.keys())
    while True:
        mod_size = input(f"选择Whisper模型（默认 small，推荐）：").strip().lower()
        if not mod_size:
            mod_size = 'small'
            print(f"✓ 使用推荐模型：{mod_size} ({AVAILABLE_WHISPER_MODELS[mod_size]['description']})\n")
            break
        if mod_size in VALID_MODELS:
            print(f"✓ 已选择模型：{mod_size} ({AVAILABLE_WHISPER_MODELS[mod_size]['description']})\n")
            break
        else:
            print(f"❌ 无效的模型名称，请从以上列表中选择\n")


    use_enhanced = False
    use_hubert = False

    if ENHANCED_MODE_AVAILABLE or HUBERT_MODE_AVAILABLE:
        print("\n转录模式选择：")
        print("1. 标准模式（快速，推荐日常使用）")
        print("   - Whisper + ffmpeg预处理")
        print("   - 速度最快，无需额外下载")
        print("   - 对法语支持良好")
        print()

        if ENHANCED_MODE_AVAILABLE:
            print("2. 增强模式（降噪优化）")
            print("   - librosa + 降噪 + Whisper")
            print("   - 适合嘈杂环境录音")
            print("   - 已测试通过（416字符）✓")
            print()

        if HUBERT_MODE_AVAILABLE:
            print("3. HuBERT模式（最高精度，法语专用）✨")
            print("   - Wav2Vec2-CTC法语模型")
            print("   - 语义特征提取")
            print("   - 专门针对法语优化")
            print("   - 已测试通过（420字符）✓")
            print()

        available_modes = ['1']
        if ENHANCED_MODE_AVAILABLE:
            available_modes.append('2')
        if HUBERT_MODE_AVAILABLE:
            available_modes.append('3')

        mode_choice = input(f"请选择模式（默认1，可选{'、'.join(available_modes)}）：").strip()

        if mode_choice == '3' and HUBERT_MODE_AVAILABLE:
            use_hubert = True
            print("✓ 已选择HuBERT模式（法语专用）\n")
        elif mode_choice == '2' and ENHANCED_MODE_AVAILABLE:
            use_enhanced = True
            print("✓ 已选择增强模式（降噪优化）\n")
        else:
            use_enhanced = False
            use_hubert = False
            print("✓ 已选择标准模式\n")
    else:
        print("\n⚠️  高级模式不可用，使用标准模式")
        print("   要启用，请安装：pip install librosa noisereduce transformers\n")

    print(f"🔄 正在转录音频（源语言：{'中文' if source_lang=='zh' else '法语'}）...")
    json_filepath = transcribe_audio_to_json(
        str(audio_file),
        model_size=mod_size,
        output_dir="数据/Jason转录数据",
        use_enhanced_mode=use_enhanced if ENHANCED_MODE_AVAILABLE else False,
        use_hubert_mode=use_hubert if HUBERT_MODE_AVAILABLE else False,
        source_language=source_lang
    )

    if not json_filepath:
        print("❌ 转录失败")
        return


    json_path = Path(json_filepath)
    with open(json_path, 'r', encoding='utf-8') as f:
        transcript_data = json.load(f)
        original_text = transcript_data['transcript']

    print(f"\n✅ 转录完成：{original_text[:100]}{'...' if len(original_text) > 100 else ''}")
    print(f"   JSON文件：{json_path.name}")
    print(f"   检测到文本长度：{len(original_text)} 字符")


    print("=" * 70)
    print("\n【步骤 3/5】翻译文本")
    print("-" * 70)
    print("🔄 正在翻译...")

    success, translated_text, trans_chars, error_msg = translate_text(original_text, direction)

    if not success:
        print(f"❌ 翻译失败：{error_msg}")
        return

    print(f"\n✅ 翻译完成：{translated_text[:100]}{'...' if len(translated_text) > 100 else ''}")
    print(f"   使用字符数：{trans_chars}")


    print("\n" + "=" * 70)
    print("\n【步骤 4/5】AI润色优化")
    print("-" * 70)
    print("🔄 正在使用通义千问AI润色译文...")

    polish_success, polished_text, qwen_tokens, polish_error = polish_text(translated_text, target_lang)

    if not polish_success:
        print(f"⚠️  润色失败：{polish_error}")
        print("将使用原始译文继续...")
        final_text = translated_text
        qwen_tokens = 0
    else:
        print(f"\n✅ 润色完成：{polished_text[:100]}{'...' if len(polished_text) > 100 else ''}")
        print(f"   使用Token数：{qwen_tokens}")
        final_text = polished_text


    print("\n" + "=" * 70)
    print("\n【步骤 5/5】文本转语音")
    print("-" * 70)


    output_dir = Path(__file__).parent.parent / "数据" / "翻译结果"
    output_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_filename = f"translated_{timestamp}.mp3"
    output_path = output_dir / output_filename

    print(f"🔄 正在生成语音（{target_lang}）...")
    print(f"   （使用{'润色后' if polish_success else '原始'}译文）")

    tts_success, tts_chars = text_to_speech(
        final_text,
        str(output_path),
        language=target_lang
    )

    if not tts_success:
        print("❌ 语音生成失败")
        return

    print(f"   使用字符数：{tts_chars}")


    print("\n" + "=" * 70)
    print("\n【保存历史记录】")
    print("-" * 70)

    total_chars = trans_chars + tts_chars

    history_data = {
        "时间戳": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "音频文件": str(audio_file),
        "Whisper模型": mod_size,
        "翻译方向": direction,
        "原文": original_text,
        "原始译文": translated_text,
        "AI润色结果": {
            "是否润色": polish_success,
            "润色后译文": polished_text if polish_success else None,
            "最终使用": final_text
        },
        "输出音频": str(output_path),
        "字符使用情况": {
            "翻译消耗": trans_chars,
            "TTS消耗": tts_chars,
            "总计": total_chars
        },
        "通义千问Token使用": qwen_tokens
    }

    history_file = save_translation_history(history_data)
    print(f"✅ 历史记录已保存：{history_file}")


    new_total_chars = get_total_chars_used()
    new_chars_remaining = XFYUN_TOKEN_LIMIT - new_total_chars
    new_chars_percentage = (new_total_chars / XFYUN_TOKEN_LIMIT) * 100

    new_qwen_tokens = get_total_qwen_tokens_used()
    new_qwen_remaining = QWEN_TOKEN_LIMIT - new_qwen_tokens
    new_qwen_percentage = (new_qwen_tokens / QWEN_TOKEN_LIMIT) * 100

    print(f"\n📊 更新后的使用统计：")
    print(f"\n   【讯飞API】")
    print(f"   累计使用：{new_total_chars}/{XFYUN_TOKEN_LIMIT} ({new_chars_percentage:.1f}%)")
    print(f"   剩余可用：{new_chars_remaining} 字符")

    print(f"\n   【通义千问】")
    print(f"   累计使用：{new_qwen_tokens}/{QWEN_TOKEN_LIMIT} ({new_qwen_percentage:.1f}%)")
    print(f"   剩余可用：{new_qwen_remaining} tokens")

    if new_chars_remaining < 500:
        print(f"\n   ⚠️  警告：讯飞API剩余字符不足！")
    elif new_chars_remaining < 2000:
        print(f"\n   ⚡ 提示：讯飞API字符即将用完")

    if new_qwen_remaining < 500:
        print(f"   ⚠️  警告：通义千问剩余Token不足！")
    elif new_qwen_remaining < 1000:
        print(f"   ⚡ 提示：通义千问Token即将用完")


    print("\n" + "=" * 70)
    print("✅ 所有步骤完成！")
    print("=" * 70)
    print(f"\n📁 输出音频：{output_path}")
    print(f"📄 历史记录：{history_file}")
    print(f"\n💬 本次消耗：")
    print(f"   讯飞API：{total_chars} 字符（翻译{trans_chars} + TTS{tts_chars}）")
    print(f"   通义千问：{qwen_tokens} tokens（AI润色）")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断操作")
    except Exception as e:
        print(f"\n\n❌ 程序异常：{e}")
        import traceback
        traceback.print_exc()
